///////////////////////////////////////////////
// This file is part of NoLifeStory.         //
// Please see Global.h for more information. //
///////////////////////////////////////////////
#include "../Global.h"

#ifdef ONLINE
#include "Crypto++/aes.h"
#include "Crypto++/modes.h"

NetworkSession *ns = nullptr;

NetworkSession::NetworkSession(string IP, unsigned short port) {
	header = true;
	use_encryption = false;

	// Some WSA startup
	WSADATA info;
	if (WSAStartup(MAKEWORD(2, 0), &info)) {
		throw "Could not start up WSA...";
	}

	sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (sock == INVALID_SOCKET) {
		throw "Could not start up connection: " + WSAGetLastError();
	}

	// Set up the connection thing
	sockaddr_in address;
	address.sin_family = AF_INET;
	address.sin_addr.S_un.S_addr = inet_addr(IP.c_str());
	address.sin_port = htons(port);
	connect(sock, (sockaddr *)&address, sizeof(address));

	int error = WSAGetLastError();
	if (error == 10060) {
		string errorstr = "Could not connect to the MapleStory server. Seems that the server is offline. WSAGetLastError(): ";
		errorstr += tostring(error);
		throw exception(errorstr.c_str());
	} else if (error != 0) {
		string errorstr = "Could not connect to the MapleStory server. Unknown error. WSAGetLastError(): ";
		errorstr += tostring(error);
		throw exception(errorstr.c_str());
	}

	unsigned long BlockingMode = 0;
	ioctlsocket(sock, FIONBIO, &BlockingMode);

	printf("Successfully established connection to %s:%d!\n", IP.c_str(), port);
}

void NetworkSession::SendPacket(const PacketWriter &packet, bool encrypt) {
	cout << "Sending packet: " << packet << endl;
	SendData(packet.GetBuffer(), packet.GetSize(), encrypt);
}

void NetworkSession::SendData(const unsigned char* buffer, int length, bool encrypt) {
	size_t realLength = encrypt ? length + HEADER_LEN : length;
	unsigned char *buf = new unsigned char[realLength];

	if (encrypt) {
		memcpy(buf + HEADER_LEN, buffer, length);

		CreateHeader(buf, (short)length);
		EncryptPacket(buf + HEADER_LEN, length);
		NextIV(IV_SEND);
	}
	else {
		memcpy(buf, buffer, length);
	}

	int s = send(sock, reinterpret_cast<const char*>(buf), realLength, 0x0);
	if (s == SOCKET_ERROR) {
		cout << "Socket error while sending packet data: " << WSAGetLastError() << endl;
	}
}

DWORD _stdcall NetworkSession::CreateReceiveThread(void *param) {
	NetworkSession *self = (NetworkSession*)param;
	return self->ReceiveAndHandlePackets();
}

DWORD NetworkSession::ReceiveAndHandlePackets() {
	array<unsigned char, BUFFER_LEN> buffer;
	int nextlen = HEADER_NO_ENC_LEN;
	while (true) {
		try {
			if (!use_encryption) {
				if (header) {
					ReceiveHeader(buffer.data());
					PacketReader pr = PacketReader(buffer.data(), nextlen);
					nextlen = pr.Read<short>();
					header = false;
				}
				else {
					ReceiveData(buffer.data(), nextlen);
					PacketReader pr(buffer.data(), nextlen);
					mapleversion = pr.Read<short>();
					string patchloc = pr.ReadString();
					unsigned int sendiv = pr.Read<int>();
					unsigned int recviv = pr.Read<int>();
					SetSendIV(reinterpret_cast<unsigned char*>(&sendiv));
					SetRecvIV(reinterpret_cast<unsigned char*>(&recviv));

					mapletype = pr.Read<uint8_t>();
					
					nextlen = HEADER_LEN;
					header = true;
					use_encryption = true;
					cout << "Connection to MapleStory server established!" << endl;
					cout << "Version: " << mapleversion << " | Patch location: " << patchloc << " | Locale: " << (int)mapletype << endl;
					InitializeHeaderHandlers();
				}
			}
			else {
				if (header) {
					ReceiveHeader(buffer.data());
					nextlen = GetLengthFromHeader(buffer.data());
					header = false;
				}
				else {
					ReceiveData(buffer.data(), nextlen);
					DecryptPacket(buffer.data(), nextlen);
					PacketReader pr(buffer.data(), nextlen);
					cout << "Got packet: " << pr << endl;
					uint16_t h = pr.Read<uint16_t>();
					if (headers[h] != nullptr) {
						headers[h](pr);
					} else {
						cout << "[RECV] Unknown packet: " << pr << endl;
					}
					nextlen = HEADER_LEN;
					header = true;
				}
			}
		}
		catch (exception &ex) {
			cout << "Exception while receiving data: " << ex.what() << endl;
			return WSAGetLastError();
		}
	}
	cout << "Exited receive function." << endl;
	return 0;
}

void NetworkSession::InitializeHeaderHandlers() {
	if (mapletype == 0x08) {
		if (mapleversion >= 75) {
			headers[0x00] = &ServerHeaderHandler::HandleLogonResult;
			headers[0x06] = &ServerHeaderHandler::HandlePINResult;
			headers[0x11] = &ServerHeaderHandler::HandlePing;
			headers[0x0A] = &Worlds::LoadWorld;
		}
	}
	
	if (headers.size() == 0) {
		throw exception("There are no headers handled for this version (yet)");
	}
}

unsigned short NetworkSession::GetLengthFromHeader(unsigned char* header) {
	int length = (int)header[0] |
					(int)(header[1] << 8) |
					(int)(header[2] << 16) |
					(int)(header[3] << 24);
	length = (length >> 16) ^ (length & 0xFFFF);
	return (unsigned short)length;
}

void NetworkSession::ReceiveHeader(unsigned char* buffer) {
	ReceiveData(buffer, use_encryption ? HEADER_LEN : HEADER_NO_ENC_LEN);
}

void NetworkSession::ReceiveData(unsigned char* buffer, int length) {
	int received = 0;
	while (true) {
		if (received == length) break;
		int r = recv(sock, (char *)(buffer + received), length - received, 0x00);
		if (r <= 0) {
			throw exception("Received less or equal to 0 bytes. See WSAGetLastError result.");
		}
		received += r;
	}
	return;
}

// Crypto.
const uint8_t AesKeySize = 32;
const uint8_t AesKey[AesKeySize] = {
	0x13, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x06, 0x00, 0x00, 0x00, 0xB4, 0x00, 0x00, 0x00,
	0x1B, 0x00, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x33, 0x00, 0x00, 0x00, 0x52, 0x00, 0x00, 0x00
};
const uint8_t IvShiftKey[256] = {
	0xEC, 0x3F, 0x77, 0xA4, 0x45, 0xD0, 0x71, 0xBF, 0xB7, 0x98, 0x20, 0xFC, 0x4B, 0xE9, 0xB3, 0xE1,
	0x5C, 0x22, 0xF7, 0x0C, 0x44, 0x1B, 0x81, 0xBD, 0x63, 0x8D, 0xD4, 0xC3, 0xF2, 0x10, 0x19, 0xE0,
	0xFB, 0xA1, 0x6E, 0x66, 0xEA, 0xAE, 0xD6, 0xCE, 0x06, 0x18, 0x4E, 0xEB, 0x78, 0x95, 0xDB, 0xBA,
	0xB6, 0x42, 0x7A, 0x2A, 0x83, 0x0B, 0x54, 0x67, 0x6D, 0xE8, 0x65, 0xE7, 0x2F, 0x07, 0xF3, 0xAA,
	0x27, 0x7B, 0x85, 0xB0, 0x26, 0xFD, 0x8B, 0xA9, 0xFA, 0xBE, 0xA8, 0xD7, 0xCB, 0xCC, 0x92, 0xDA,
	0xF9, 0x93, 0x60, 0x2D, 0xDD, 0xD2, 0xA2, 0x9B, 0x39, 0x5F, 0x82, 0x21, 0x4C, 0x69, 0xF8, 0x31,
	0x87, 0xEE, 0x8E, 0xAD, 0x8C, 0x6A, 0xBC, 0xB5, 0x6B, 0x59, 0x13, 0xF1, 0x04, 0x00, 0xF6, 0x5A,
	0x35, 0x79, 0x48, 0x8F, 0x15, 0xCD, 0x97, 0x57, 0x12, 0x3E, 0x37, 0xFF, 0x9D, 0x4F, 0x51, 0xF5,
	0xA3, 0x70, 0xBB, 0x14, 0x75, 0xC2, 0xB8, 0x72, 0xC0, 0xED, 0x7D, 0x68, 0xC9, 0x2E, 0x0D, 0x62,
	0x46, 0x17, 0x11, 0x4D, 0x6C, 0xC4, 0x7E, 0x53, 0xC1, 0x25, 0xC7, 0x9A, 0x1C, 0x88, 0x58, 0x2C,
	0x89, 0xDC, 0x02, 0x64, 0x40, 0x01, 0x5D, 0x38, 0xA5, 0xE2, 0xAF, 0x55, 0xD5, 0xEF, 0x1A, 0x7C,
	0xA7, 0x5B, 0xA6, 0x6F, 0x86, 0x9F, 0x73, 0xE6, 0x0A, 0xDE, 0x2B, 0x99, 0x4A, 0x47, 0x9C, 0xDF,
	0x09, 0x76, 0x9E, 0x30, 0x0E, 0xE4, 0xB2, 0x94, 0xA0, 0x3B, 0x34, 0x1D, 0x28, 0x0F, 0x36, 0xE3,
	0x23, 0xB4, 0x03, 0xD8, 0x90, 0xC8, 0x3C, 0xFE, 0x5E, 0x32, 0x24, 0x50, 0x1F, 0x3A, 0x43, 0x8A,
	0x96, 0x41, 0x74, 0xAC, 0x52, 0x33, 0xF0, 0xD9, 0x29, 0x80, 0xB1, 0x16, 0xD3, 0xAB, 0x91, 0xB9,
	0x84, 0x7F, 0x61, 0x1E, 0xCF, 0xC5, 0xD1, 0x56, 0x3D, 0xCA, 0xF4, 0x05, 0xC6, 0xE5, 0x08, 0x49
};

void NetworkSession::DecryptPacket(unsigned char* data, int length) {
	int pos = 0;
	uint8_t first = 1; // Needed for the 4 bytes skip.

	// Initialize Crypto++ decryption class
	CryptoPP::OFB_Mode<CryptoPP::AES>::Decryption crypto;
	
	while (length > pos) {
		crypto.SetKeyWithIV(AesKey, AesKeySize, IV_RECV);
		if (length>(pos + 1460 - first * 4)) {
			crypto.ProcessData(data + pos, data + pos, 1460 - first * 4);
		}
		else {
			crypto.ProcessData(data + pos, data + pos, length - pos);
		}
		pos += 1460 - first * 4;
		if (first) first = 0;
	}

	NextIV(IV_RECV);

	int j;
	uint8_t a, b, c;
	for (auto i = 0; i < 3; i++) {
		a = 0;
		b = 0;
		for (j = length; j > 0; j--) {
			c = data[j - 1];
			c = rol(c, 3);
			c = c ^ 0x13;
			a = c;
			c = c ^ b;
			c = (uint8_t)(c - j); // Guess this is supposed to be right?
			c = ror(c, 4);
			b = a;
			data[j - 1] = c;
		}
		a = 0;
		b = 0;
		for (j = length; j > 0; j--) {
			c = data[length - j];
			c = c - 0x48;
			c = c ^ 0xFF;
			c = rol(c, j);
			a = c;
			c = c ^ b;
			c = (uint8_t)(c - j); // Guess this is supposed to be right?
			c = ror(c, 3);
			b = a;
			data[length - j] = c;
		}
	}
}

void NetworkSession::EncryptPacket(unsigned char* data, int length) {
	int j;
	uint8_t a, b, c;
	for (auto i = 0; i < 3; i++) {
		a = 0;
		for (j = length; j > 0; j--) {
			c = data[length - j];
			c = rol(c, 3);
			c = (uint8_t)(c + j); // Guess this is supposed to be right?
			c = c ^ a;
			a = c;
			c = ror(a, j);
			c = c ^ 0xFF;
			c = c + 0x48;
			data[length - j] = c;
		}
		a = 0;
		for (j = length; j > 0; j--) {
			c = data[j - 1];
			c = rol(c, 4);
			c = (uint8_t)(c + j); // Guess this is supposed to be right?
			c = c ^ a;
			a = c;
			c = c ^ 0x13;
			c = ror(c, 3);
			data[j - 1] = c;
		}
	}

	int pos = 0;
	uint8_t first = 1; // Needed for the 4 bytes skip.

	// Initialize Crypto++ encryption class
	CryptoPP::OFB_Mode<CryptoPP::AES>::Encryption crypto;

	while (length > pos) {
		crypto.SetKeyWithIV(AesKey, AesKeySize, IV_SEND);
		if (length>(pos + 1460 - first * 4)) {
			crypto.ProcessData(data + pos, data + pos, 1460 - first * 4);
		}
		else {
			crypto.ProcessData(data + pos, data + pos, length - pos);
		}
		pos += 1460 - first * 4;
		if (first) first = 0;
	}
}

void NetworkSession::NextIV(unsigned char *oldiv) {
	uint8_t newIV[4] = { 0xF2, 0x53, 0x50, 0xC6 };

	for (auto i = 0; i < 4; i++) {
		uint8_t temp1 = newIV[1];
		uint8_t temp2 = IvShiftKey[temp1];
		uint8_t temp3 = oldiv[i];
		temp2 -= temp3;
		newIV[0] += temp2;
		temp2 = newIV[2];
		temp2 ^= IvShiftKey[temp3];
		temp1 -= temp2;
		newIV[1] = temp1;
		temp1 = newIV[3];
		temp2 = temp1;
		temp1 -= newIV[0];
		temp2 = IvShiftKey[temp2];
		temp2 += temp3;
		temp2 ^= newIV[2];
		newIV[2] = temp2;
		temp1 += IvShiftKey[temp3];
		newIV[3] = temp1;
		unsigned int result1 = (unsigned int)newIV[0] | ((unsigned int)newIV[1] << 8) | ((unsigned int)newIV[2] << 16) | ((unsigned int)newIV[3] << 24);
		unsigned int result2 = result1 >> 0x1D;
		result1 <<= 3;
		result2 |= result1;
		newIV[0] = (uint8_t)(result2 & 0xFF);
		newIV[1] = (uint8_t)((result2 >> 8) & 0xFF);
		newIV[2] = (uint8_t)((result2 >> 16) & 0xFF);
		newIV[3] = (uint8_t)((result2 >> 24) & 0xFF);
	}
	SetIV(newIV, oldiv);
}

void NetworkSession::CreateHeader(unsigned char* header, short length) {
	unsigned short a = (IV_SEND[3] * 0x100) + IV_SEND[2];
	a ^= mapleversion;
	unsigned short b = a ^ length;
	header[0] = a % 0x100;
	header[1] = a / 0x100;
	header[2] = b % 0x100;
	header[3] = b / 0x100;
}
#endif